from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
import time
import base64
from PIL import Image
import io
import os

def capture_scroll_screenshot_and_save_pdf(url, output_pdf="output.pdf", scroll_step=500, max_scrolls=20, timeout=10):
    """
    滚动截屏并拼接为 PDF
    
    :param url: 目标网页 URL
    :param output_pdf: 输出 PDF 路径（默认：output.pdf）
    :param scroll_step: 每次滚动的像素步长（默认：500px，可根据窗口高度调整）
    :param max_scrolls: 最大滚动次数（防止无限循环，默认：20）
    :param timeout: 内容加载超时时间（秒，默认：10）
    """
    # 配置 Chrome 选项（无头模式 + 截图优化）
    chrome_options = Options()
    chrome_options.add_argument("--headless=new")  # 无头模式（Chrome 112+）
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--window-size=1920,1080")  # 窗口尺寸（影响截图宽度和渲染）
    chrome_options.add_argument("--hide-scrollbars")  # 隐藏滚动条（避免截图出现黑边）

    # 初始化浏览器驱动
    driver = webdriver.Chrome(options=chrome_options)

    try:
        # 导航到目标页面并等待初始加载
        driver.get(url)
        WebDriverWait(driver, timeout).until(
            EC.presence_of_element_located(("tag name", "body"))  # 等待页面主体加载
        )

        # 获取初始窗口尺寸（用于计算截图位置）
        window_height = driver.execute_script("return window.innerHeight")
        scroll_position = 0  # 当前滚动位置
        screenshots = []  # 存储所有分段截图的二进制数据

        # 模拟滚动并截图
        for _ in range(max_scrolls):
            # 滚动到当前位置（每次滚动 step 像素）
            driver.execute_script(f"window.scrollTo(0, {scroll_position});")
            scroll_position += scroll_step

            print(f'滚动到{scroll_position}')

            # 等待页面稳定（避免滚动动画未完成导致截图模糊）
            time.sleep(5)

            # 截取当前视口内容（返回 Base64 编码的 PNG 图片）
            screenshot_base64 = driver.get_screenshot_as_base64()
            screenshot_bytes = base64.b64decode(screenshot_base64)
            img = Image.open(io.BytesIO(screenshot_bytes))

            # 裁剪截图（去除可能的底部空白，可选）
            # 若页面高度不足窗口高度，裁剪到实际内容高度
            actual_height = min(window_height, img.height)
            img = img.crop((0, 0, img.width, actual_height))

            screenshots.append(img)
            print(f'捕获到第{len(screenshots)}幅画面')

            # 检测是否已滚动到页面底部（通过比较当前滚动位置和总高度）
            total_height = driver.execute_script("return document.documentElement.scrollHeight")
            if scroll_position >= total_height - window_height:
                break  # 已到底部，退出循环

        # 拼接所有截图为一张长图
        if not screenshots:
            raise ValueError("未获取到任何截图，请检查页面加载或滚动逻辑")

        # 计算总高度和单张宽度
        total_width = screenshots[0].width
        total_height = sum(img.height for img in screenshots)

        print(f'合并总画面{total_width}x{total_height}')

        # 创建空白长图（模式与截图一致，通常为 RGB）
        combined_img = Image.new("RGB", (total_width, total_height))

        # 粘贴每张截图到正确位置
        y_offset = 0
        for img in screenshots:
            combined_img.paste(img, (0, y_offset))
            y_offset += img.height

        # 保存长图为 PDF（Pillow 直接支持）
        combined_img.save(output_pdf, "PDF", resolution=100.0)
        print(f"PDF 成功保存至：{output_pdf}")

    except Exception as e:
        print(f"生成失败，错误信息：{str(e)}")
    finally:
        driver.quit()  # 确保关闭浏览器

# 使用示例
if __name__ == "__main__":
    target_url = "http://114.251.111.103:18080/kfs/file/read/96fff55c0628d9342c2cacbd786c0bce"  # 替换为实际长页面 URL（如维基百科长文）
    capture_scroll_screenshot_and_save_pdf(
        url=target_url,
        output_pdf="long_page.pdf",
        scroll_step=800,  # 根据窗口高度调整（如窗口高 1080，可设 800 避免重叠）
        max_scrolls=150,   # 根据页面长度调整最大滚动次数
        timeout=120        # 根据网络速度调整加载超时时间
    )